"""
1.	利用Pytorch深度学习框架，建立由卷积神经网络Resnet34构造的目标检测YOLOv1主干模型，和后端网络，
并根据给定的YOLOv1算法的loss损失函数公式，按下面要求完成相应代码。（40分）
"""
import torch as pt
import torchvision as ptv
import matplotlib.pyplot as plt
import numpy as np

pt.manual_seed(1)
np.random.seed(1)


def sep(label=''):
    print('-' * 32, label, '-' * 32, sep='')


sep('Yolo v2 model')


def ConvBnRelu(in_ch, out_ch, ksize=(3, 3), stride=(1, 1), padding=1, bias=False):
    return pt.nn.Sequential(
        pt.nn.Conv2d(in_ch, out_ch, ksize, stride, padding, bias=bias),
        pt.nn.BatchNorm2d(out_ch),
        pt.nn.ReLU(),
    )

# ①	建立YOLOv1模型类结构
class YoloV1Model(pt.nn.Module):

    def __init__(self, n_cls, **kwargs):
        super().__init__(**kwargs)
        depth = 5 * 2 + n_cls

        # ②	导入Resnet34主干网络作为与训练模型
        resnet = ptv.models.resnet34(pretrained=True)
        out_ch = resnet.fc.in_features
        self.resnet = pt.nn.Sequential(*(list(resnet.children())[:-2]))

        # ③	搭建后端网络模型
        self.customer = pt.nn.Sequential(
            ConvBnRelu(out_ch, 1024),
            ConvBnRelu(1024, 1024, (3, 3), (2, 2)),
            ConvBnRelu(1024, 1024),
            ConvBnRelu(1024, 1024),
            ConvBnRelu(1024, depth),
        )

    def forward(self, x):
        x = self.resnet(x)
        x = self.customer(x)
        return x


# ④	打印输出YOLOv1整体模型维度结构
n_cls = 20
depth = 5 * 2 + n_cls
model = YoloV1Model(n_cls)
print(model)

# ⑤	参考YOLOv1算法的loss损失函数计算公式（参见题末）
sep('Yolo v2 loss')


def get_area(box):
    return (box[2] - box[0] + 1) * (box[3] - box[1] + 1)


def get_iou(box, gt):
    max_x1 = pt.maximum(box[0], gt[0])
    min_x2 = pt.minimum(box[2], gt[2])
    max_y1 = pt.maximum(box[1], gt[1])
    min_y2 = pt.minimum(box[3], gt[3])
    xx = pt.maximum(pt.Tensor([0.]), min_x2 - max_x1 + 1)
    yy = pt.maximum(pt.Tensor([0.]), min_y2 - max_y1 + 1)
    I = xx * yy
    U = get_area(box) + get_area(gt) - I
    iou = I / (U + 1e-20)
    return iou


def yolo_v1_loss(y_true, pred):
    lambda_box, lambda_nonobj = 5., 0.5
    loss_conf_obj, loss_conf_nonobj = 0., 0.
    loss_box, loss_cls = 0., 0.
    B, CH, H, W = y_true.size()
    for sample_id in range(B):
        for row_id in range(H):
            for col_id in range(W):
                # ⑥	编写目标和非目标置信度代价代码
                if pt.isclose(y_true[sample_id, 4, row_id, col_id], pt.Tensor([0.])):
                    loss_conf_nonobj += lambda_nonobj * (pred[sample_id, [4, 9], row_id, col_id] ** 2).sum()
                else:
                    # calculate IOU of the boxes between ground truth
                    box01 = (
                        col_id + pred[sample_id, 0, row_id, col_id] - W * pred[sample_id, 2, row_id, col_id] * 0.5,
                        row_id + pred[sample_id, 1, row_id, col_id] - H * pred[sample_id, 3, row_id, col_id] * 0.5,
                        col_id + pred[sample_id, 0, row_id, col_id] + W * pred[sample_id, 2, row_id, col_id] * 0.5,
                        row_id + pred[sample_id, 1, row_id, col_id] + H * pred[sample_id, 3, row_id, col_id] * 0.5,
                    )
                    box02 = (
                        col_id + pred[sample_id, 5, row_id, col_id] - W * pred[sample_id, 7, row_id, col_id] * 0.5,
                        row_id + pred[sample_id, 6, row_id, col_id] - H * pred[sample_id, 8, row_id, col_id] * 0.5,
                        col_id + pred[sample_id, 5, row_id, col_id] + W * pred[sample_id, 7, row_id, col_id] * 0.5,
                        row_id + pred[sample_id, 6, row_id, col_id] + H * pred[sample_id, 8, row_id, col_id] * 0.5,
                    )
                    gt = (
                        col_id + y_true[sample_id, 0, row_id, col_id] - W * y_true[sample_id, 2, row_id, col_id] * 0.5,
                        row_id + y_true[sample_id, 1, row_id, col_id] - H * y_true[sample_id, 3, row_id, col_id] * 0.5,
                        col_id + y_true[sample_id, 0, row_id, col_id] + W * y_true[sample_id, 2, row_id, col_id] * 0.5,
                        row_id + y_true[sample_id, 1, row_id, col_id] + H * y_true[sample_id, 3, row_id, col_id] * 0.5,
                    )

                    iou01 = get_iou(box01, gt)
                    iou02 = get_iou(box02, gt)

                    # ⑦	编写目标边框代价代码
                    if iou01 > iou02:
                        loss_conf_obj += ((pred[sample_id, 4, row_id, col_id] - iou01) ** 2).sum()
                        loss_conf_nonobj += lambda_nonobj * ((pred[sample_id, 9, row_id, col_id] - iou02) ** 2).sum()
                        loss_box += lambda_box * (
                            ((pred[sample_id, 0:2, row_id, col_id] - y_true[sample_id, 0:2, row_id, col_id]) ** 2).sum()
                            + ((pred[sample_id, 2:4, row_id, col_id].sqrt() - y_true[sample_id, 2:4, row_id, col_id].sqrt()) ** 2).sum()
                        )
                    else:
                        loss_conf_obj += ((pred[sample_id, 9, row_id, col_id] - iou02) ** 2).sum()
                        loss_conf_nonobj += lambda_nonobj * ((pred[sample_id, 4, row_id, col_id] - iou01) ** 2).sum()
                        loss_box += lambda_box * (
                                ((pred[sample_id, 5:7, row_id, col_id] - y_true[sample_id, 5:7, row_id, col_id]) ** 2).sum()
                                + ((pred[sample_id, 7:9, row_id, col_id].sqrt() - y_true[sample_id, 7:9, row_id, col_id].sqrt()) ** 2).sum()
                        )

                    # ⑧	编写目标分类代价代码
                    loss_cls += ((y_true[sample_id, 5:, row_id, col_id] - pred[sample_id, 5:, row_id, col_id]) ** 2).sum()
    loss = loss_conf_nonobj + loss_conf_obj + loss_box + loss_cls
    loss /= B
    return loss


# Try to calculate a loss for YOLO v1
x = pt.zeros(4, 3, 448, 448, dtype=pt.float32)
pred = model(x)
print('pred.size', pred.size())
y_true = pt.zeros(4, depth, 7, 7, dtype=pt.float32)
y_true[0, [4, 9], 6, 6] = 1.
y_true[1, [4, 9], 0, 0] = 1.
y_true[2, [4, 9], 0, 6] = 1.
y_true[3, [4, 9], 4, 4] = 1.
loss = yolo_v1_loss(y_true, pred).detach().cpu().item()
print(f'Loss = {loss}')

# ⑨	根据给定的五个VOC“锚框”Anchors的宽高数据，和目标真实框GT数据，分别计算GT与五个Anchors的交并比IOU，并输出他们的值
# Anchors: (1.3221, 1.73145), (3.19275, 4.00944), (5.05587, 8.09892), (9.47112, 4.84053), (11.2364, 10.0071)
# GT：中心点(7, 7)，边框宽高(10, 5)
sep('Anchors and GT')
anchors = np.array((
    (1.3221, 1.73145), (3.19275, 4.00944), (5.05587, 8.09892), (9.47112, 4.84053), (11.2364, 10.0071)
))
gt = np.array((
    (10, 5),
))

center_val = 7
boxes = np.array((
    center_val - anchors[:, 0] / 2,
    center_val - anchors[:, 1] / 2,
    center_val + anchors[:, 0] / 2,
    center_val + anchors[:, 1] / 2,
)).transpose()
gt_box = np.array((
    center_val - gt[:, 0] / 2,
    center_val - gt[:, 1] / 2,
    center_val + gt[:, 0] / 2,
    center_val + gt[:, 1] / 2,
)).transpose()


def get_area_np(box):
    return (box[:, 2] - box[:, 0] + 1) * (box[:, 3] - box[:, 1] + 1)


def get_iou_np(box, gt):
    max_x1 = np.maximum(box[:, 0], gt[:, 0])
    min_x2 = np.minimum(box[:, 2], gt[:, 2])
    max_y1 = np.maximum(box[:, 1], gt[:, 1])
    min_y2 = np.minimum(box[:, 3], gt[:, 3])
    xx = np.maximum(0., min_x2 - max_x1 + 1)
    yy = np.maximum(0., min_y2 - max_y1 + 1)
    I = xx * yy
    U = get_area_np(box) + get_area_np(gt) - I
    iou = I / (U + 1e-20)
    return iou
iou = get_iou_np(boxes, gt_box)
print('IOU:', iou)


def draw_boxes(boxes, orders, color, label=None):
    x1 = boxes[orders, 0]
    y1 = boxes[orders, 1]
    x2 = boxes[orders, 2]
    y2 = boxes[orders, 3]
    plt.plot([x1, x2], [y1, y1], color=color)
    plt.plot([x1, x2], [y2, y2], color=color)
    plt.plot([x1, x1], [y1, y2], color=color)
    plt.plot([x2, x2], [y1, y2], color=color)
    if label is None:
        show_idx = True
    else:
        show_idx = False
    for o in orders:
        if show_idx:
            label = f'#{o}'
        plt.annotate(label, xy=[boxes[o, 0], boxes[o, 1]], color=color)


# ⑩	选择与真实框GT最大交并比的“锚框”Anchor，图像显示该“锚框”和真实框GT
orders = iou.argsort()[::-1]
best = orders[0:1]
other = orders[1:]
plt.figure(figsize=[6, 6])
draw_boxes(boxes, other, 'black')
draw_boxes(gt_box, [0], 'red', 'GT')
draw_boxes(boxes, best, 'blue', 'best')

# Finally
plt.ylim(plt.ylim()[::-1])
plt.show()
